# encoding:utf-8
#
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from npu_bridge.npu_init import *
# -*- coding:utf-8 -*-
import tensorflow as tf

from .network import Network
from ..fast_rcnn.config import cfg


class VGGnet_train(Network):
    def __init__(self, trainable=True):
        self.inputs = []
        self.data = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='data')
        util.set_graph_exec_config(self.data, dynamic_input=True, dynamic_inputs_shape_range="data:[1,1~1024,1~1024,3]")
        # self.im_info = tf.placeholder(tf.float32, shape=[None, 3], name='im_info')
        # self.gt_boxes = tf.placeholder(tf.float32, shape=[None, 5], name='gt_boxes')
        # self.gt_ishard = tf.placeholder(tf.int32, shape=[None], name='gt_ishard')
        # self.dontcare_areas = tf.placeholder(tf.float32, shape=[None, 4], name='dontcare_areas')
        self.keep_prob = tf.placeholder(tf.float32)

        self.rpn_labels = tf.placeholder(tf.int32, name='rpn_labels')
        print('**'*100, type(self.rpn_labels.name))
        self.rpn_bbox_targets = tf.placeholder(tf.float32, name='rpn_bbox_targets')
        self.rpn_bbox_inside_weights = tf.placeholder(tf.float32, name='rpn_bbox_inside_weights')
        self.rpn_bbox_outside_weights = tf.placeholder(tf.float32, name='rpn_bbox_outside_weights')
        self.layers = dict({'data': self.data, \
                            'rpn_labels':self.rpn_labels,\
                            'rpn_bbox_targets':self.rpn_bbox_targets,'rpn_bbox_inside_weights':self.rpn_bbox_inside_weights,\
                            'rpn_bbox_outside_weights':self.rpn_bbox_outside_weights})
        self.trainable = trainable

        self.setup()

    def setup(self):
        n_classes = cfg.NCLASSES
        anchor_scales = cfg.ANCHOR_SCALES
        _feat_stride = [16, ]
        # net frame
        (self.feed('data')
         .conv(3, 3, 64, 1, 1, name='conv1_1')
         .conv(3, 3, 64, 1, 1, name='conv1_2')
         .max_pool(2, 2, 2, 2, padding='VALID', name='pool1')
         .conv(3, 3, 128, 1, 1, name='conv2_1')
         .conv(3, 3, 128, 1, 1, name='conv2_2')
         .max_pool(2, 2, 2, 2, padding='VALID', name='pool2')
         .conv(3, 3, 256, 1, 1, name='conv3_1')
         .conv(3, 3, 256, 1, 1, name='conv3_2')
         .conv(3, 3, 256, 1, 1, name='conv3_3')
         .max_pool(2, 2, 2, 2, padding='VALID', name='pool3')
         .conv(3, 3, 512, 1, 1, name='conv4_1')
         .conv(3, 3, 512, 1, 1, name='conv4_2')
         .conv(3, 3, 512, 1, 1, name='conv4_3')
         .max_pool(2, 2, 2, 2, padding='VALID', name='pool4')
         .conv(3, 3, 512, 1, 1, name='conv5_1')
         .conv(3, 3, 512, 1, 1, name='conv5_2')
         .conv(3, 3, 512, 1, 1, name='conv5_3'))
        # ========= RPN ============
        (self.feed('conv5_3')
         .conv(3, 3, 512, 1, 1, name='rpn_conv/3x3'))
        
        (self.feed('rpn_conv/3x3').Bilstm(512, 128, 512, name='lstm_o'))
        (self.feed('lstm_o').lstm_fc(512, len(anchor_scales) * 10 * 4, name='rpn_bbox_pred'))
        (self.feed('lstm_o').lstm_fc(512, len(anchor_scales) * 10 * 2, name='rpn_cls_score'))

        # generating training labels on the fly
        # output: rpn_labels(HxWxA, 2) rpn_bbox_targets(HxWxA, 4) rpn_bbox_inside_weights rpn_bbox_outside_weights
        # 给每个anchor上标签，并计算真值（也是delta的形式），以及内部权重和外部权重
        # (self.feed('rpn_cls_score', 'gt_boxes', 'gt_ishard', 'dontcare_areas', 'im_info')
        #  .anchor_target_layer(_feat_stride, anchor_scales, name='rpn-data'))

        # shape is (1, H, W, Ax2) -> (1, H, WxA, 2)
        # 给之前得到的score进行softmax，得到0-1之间的得分
        (self.feed('rpn_cls_score')
         .spatial_reshape_layer(2, name='rpn_cls_score_reshape')
         .spatial_softmax(name='rpn_cls_prob'))
